//! 连接管理模块
//! 负责管理 WebSocket 和 SSE 连接的生命周期

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, RwLock};
use uuid::Uuid;

use crate::models::token_model::TokenModel;
use super::message::HubResponse;

/// 连接标识符
pub type ConnectionId = Uuid;

/// 连接类型
#[derive(Debug, Clone, PartialEq)]
pub enum ConnectionType {
    WebSocket,
    SSE,
}

/// 连接状态
#[derive(Debug, Clone)]
pub struct ConnectionState {
    pub id: ConnectionId,
    pub connection_type: ConnectionType,
    pub token: Option<TokenModel>,
    pub last_activity: std::time::Instant,
    pub is_authenticated: bool,
}

impl ConnectionState {
    pub fn new(connection_type: ConnectionType) -> Self {
        Self {
            id: Uuid::new_v4(),
            connection_type,
            token: None,
            last_activity: std::time::Instant::now(),
            is_authenticated: false,
        }
    }
    
    pub fn update_activity(&mut self) {
        self.last_activity = std::time::Instant::now();
    }
    
    pub fn authenticate(&mut self, token: TokenModel) {
        self.token = Some(token);
        self.is_authenticated = true;
        self.update_activity();
    }
    
    pub fn is_expired(&self, timeout_duration: u64) -> bool {
        self.last_activity.elapsed().as_secs() > timeout_duration
    }
}

/// 连接发送器 trait
#[async_trait::async_trait]
pub trait ConnectionSender: Send + Sync {
    /// 发送消息到连接
    async fn send(&self, message: HubResponse) -> anyhow::Result<()>;
    
    /// 获取连接ID
    fn connection_id(&self) -> ConnectionId;
}

/// WebSocket 连接发送器
pub struct WebSocketSender {
    pub connection_id: ConnectionId,
    pub sender: tokio::sync::mpsc::UnboundedSender<HubResponse>,
}

#[async_trait::async_trait]
impl ConnectionSender for WebSocketSender {
    async fn send(&self, message: HubResponse) -> anyhow::Result<()> {
        // 对于 UnboundedSender，send 方法是同步的，不需要 await
        // 但为了保持接口一致性，我们仍然使用 async
        self.sender
            .send(message)
            .map_err(|e| anyhow::anyhow!("发送消息失败: {}", e))
    }
    
    fn connection_id(&self) -> ConnectionId {
        self.connection_id
    }
}

/// SSE 连接发送器
pub struct SSESender {
    pub connection_id: ConnectionId,
    pub sender: tokio::sync::mpsc::UnboundedSender<HubResponse>,
}

#[async_trait::async_trait]
impl ConnectionSender for SSESender {
    async fn send(&self, message: HubResponse) -> anyhow::Result<()> {
        // 对于 UnboundedSender，send 方法是同步的，不需要 await
        // 但为了保持接口一致性，我们仍然使用 async
        self.sender
            .send(message)
            .map_err(|e| anyhow::anyhow!("发送消息失败: {}", e))
    }
    
    fn connection_id(&self) -> ConnectionId {
        self.connection_id
    }
}

/// 连接管理器
pub struct ConnectionManager {
    connections: Arc<RwLock<HashMap<ConnectionId, Arc<dyn ConnectionSender>>>>,
    connection_states: Arc<RwLock<HashMap<ConnectionId, ConnectionState>>>,
    message_broadcast: broadcast::Sender<HubResponse>,
    max_connections: usize,
}

impl ConnectionManager {
    pub fn new(max_connections: usize) -> Self {
        let (message_broadcast, _) = broadcast::channel(100);
        
        Self {
            connections: Arc::new(RwLock::new(HashMap::new())),
            connection_states: Arc::new(RwLock::new(HashMap::new())),
            message_broadcast,
            max_connections,
        }
    }
    
    /// 添加新连接
    pub async fn add_connection<S>(&self, sender: S, connection_type: ConnectionType) -> ConnectionId 
    where
        S: ConnectionSender + 'static,
    {
        self.add_connection_with_token(sender, connection_type, None).await
    }
    
    /// 添加带Token的新连接
    pub async fn add_connection_with_token<S>(
        &self, 
        sender: S, 
        connection_type: ConnectionType,
        token: Option<TokenModel>
    ) -> ConnectionId 
    where
        S: ConnectionSender + 'static,
    {
        let connection_id = sender.connection_id();
        
        // 检查连接数限制
        let current_count = self.connections.read().await.len();
        if current_count >= self.max_connections {
            tracing::warn!("连接数已达上限: {}/{}", current_count, self.max_connections);
        }
        
        // 添加连接发送器
        self.connections.write().await.insert(connection_id, Arc::new(sender));
        
        // 保存连接类型用于日志记录
        let connection_type_for_log = connection_type.clone();
        
        // 添加连接状态
        let mut state = ConnectionState::new(connection_type);
        if let Some(token) = token {
            state.authenticate(token);
        }
        let is_authenticated = state.is_authenticated;
        self.connection_states.write().await.insert(connection_id, state);
        
        if is_authenticated {
            tracing::info!("新连接建立: {} ({:?}) - 已认证", connection_id, connection_type_for_log);
        } else {
            tracing::info!("新连接建立: {} ({:?})", connection_id, connection_type_for_log);
        }
        connection_id
    }
    
    /// 移除连接
    pub async fn remove_connection(&self, connection_id: ConnectionId) {
        self.connections.write().await.remove(&connection_id);
        self.connection_states.write().await.remove(&connection_id);
        
        tracing::info!("连接断开: {}", connection_id);
    }
    
    /// 发送消息到指定连接
    pub async fn send_to_connection(&self, connection_id: ConnectionId, message: HubResponse) -> anyhow::Result<()> {
        let connections = self.connections.read().await;
        if let Some(sender) = connections.get(&connection_id) {
            sender.send(message).await
        } else {
            Err(anyhow::anyhow!("连接不存在: {}", connection_id))
        }
    }
    
    /// 广播消息到所有连接
    pub async fn broadcast(&self, message: HubResponse) {
        let _ = self.message_broadcast.send(message);
    }
    
    /// 获取广播接收器
    pub fn subscribe(&self) -> broadcast::Receiver<HubResponse> {
        self.message_broadcast.subscribe()
    }
    
    /// 认证连接
    pub async fn authenticate_connection(&self, connection_id: ConnectionId, token: TokenModel) -> anyhow::Result<()> {
        let mut states = self.connection_states.write().await;
        if let Some(state) = states.get_mut(&connection_id) {
            state.authenticate(token);
            tracing::info!("连接认证成功: {}", connection_id);
            Ok(())
        } else {
            Err(anyhow::anyhow!("连接不存在: {}", connection_id))
        }
    }
    
    /// 获取连接状态
    pub async fn get_connection_state(&self, connection_id: ConnectionId) -> Option<ConnectionState> {
        self.connection_states.read().await.get(&connection_id).cloned()
    }
    
    /// 更新连接活动时间
    pub async fn update_connection_activity(&self, connection_id: ConnectionId) -> anyhow::Result<()> {
        let mut states = self.connection_states.write().await;
        if let Some(state) = states.get_mut(&connection_id) {
            state.update_activity();
            Ok(())
        } else {
            Err(anyhow::anyhow!("连接不存在: {}", connection_id))
        }
    }
    
    /// 清理过期连接
    pub async fn cleanup_expired_connections(&self, timeout_duration: u64) {
        let mut states = self.connection_states.write().await;
        let mut connections = self.connections.write().await;
        
        let expired_connections: Vec<ConnectionId> = states
            .iter()
            .filter(|(_, state)| state.is_expired(timeout_duration))
            .map(|(id, _)| *id)
            .collect();
        
        for connection_id in expired_connections {
            states.remove(&connection_id);
            connections.remove(&connection_id);
            tracing::info!("清理过期连接: {}", connection_id);
        }
    }
    
    /// 获取活跃连接数
    pub async fn active_connections(&self) -> usize {
        self.connections.read().await.len()
    }
    
    /// 获取认证连接数
    pub async fn authenticated_connections(&self) -> usize {
        self.connection_states
            .read()
            .await
            .values()
            .filter(|state| state.is_authenticated)
            .count()
    }
}

/// 连接管理器包装器
#[derive(Clone)]
pub struct ConnectionManagerHandle {
    inner: Arc<ConnectionManager>,
}

impl ConnectionManagerHandle {
    pub fn new(manager: ConnectionManager) -> Self {
        Self {
            inner: Arc::new(manager),
        }
    }
    
    pub async fn add_connection<S>(&self, sender: S, connection_type: ConnectionType) -> ConnectionId 
    where
        S: ConnectionSender + 'static,
    {
        self.inner.add_connection(sender, connection_type).await
    }
    
    pub async fn add_connection_with_token<S>(
        &self, 
        sender: S, 
        connection_type: ConnectionType,
        token: Option<TokenModel>
    ) -> ConnectionId 
    where
        S: ConnectionSender + 'static,
    {
        self.inner.add_connection_with_token(sender, connection_type, token).await
    }
    
    pub async fn remove_connection(&self, connection_id: ConnectionId) {
        self.inner.remove_connection(connection_id).await;
    }
    
    pub async fn send_to_connection(&self, connection_id: ConnectionId, message: HubResponse) -> anyhow::Result<()> {
        self.inner.send_to_connection(connection_id, message).await
    }
    
    pub async fn broadcast(&self, message: HubResponse) {
        self.inner.broadcast(message).await;
    }
    
    pub fn subscribe(&self) -> broadcast::Receiver<HubResponse> {
        self.inner.subscribe()
    }
    
    pub async fn authenticate_connection(&self, connection_id: ConnectionId, token: TokenModel) -> anyhow::Result<()> {
        self.inner.authenticate_connection(connection_id, token).await
    }
    
    pub async fn get_connection_state(&self, connection_id: ConnectionId) -> Option<ConnectionState> {
        self.inner.get_connection_state(connection_id).await
    }
    
    pub async fn update_connection_activity(&self, connection_id: ConnectionId) -> anyhow::Result<()> {
        self.inner.update_connection_activity(connection_id).await
    }
    
    pub async fn cleanup_expired_connections(&self, timeout_duration: u64) {
        self.inner.cleanup_expired_connections(timeout_duration).await;
    }
    
    pub async fn active_connections(&self) -> usize {
        self.inner.active_connections().await
    }
    
    pub async fn authenticated_connections(&self) -> usize {
        self.inner.authenticated_connections().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[tokio::test]
    async fn test_connection_manager() {
        let manager = ConnectionManager::new(10);
        
        // 测试添加和移除连接
        let (tx, _) = tokio::sync::mpsc::unbounded_channel();
        let sender = WebSocketSender {
            connection_id: Uuid::new_v4(),
            sender: tx,
        };
        
        let connection_id = manager.add_connection(sender, ConnectionType::WebSocket).await;
        assert_eq!(manager.active_connections().await, 1);
        
        manager.remove_connection(connection_id).await;
        assert_eq!(manager.active_connections().await, 0);
    }
}